"""
loss_config.py
--------------
定义训练损失函数的配置，包括：
- 单任务损失（SingleLossConfig）
- 多任务损失组合（LossConfig）
- 支持内置损失函数、FocalLoss、自定义损失函数
- 每个损失可设置权重、label_smoothing、reduction
- 提供统一接口 get_combined_loss() 计算总损失
- 文件末尾包含测试代码
"""

from dataclasses import dataclass, field
from typing import Optional, Callable, List, Dict, Union
import torch
import torch.nn as nn

# -----------------------------
# 单个损失配置
# -----------------------------
@dataclass
class SingleLossConfig:
    loss_type: Optional[str] = "cross_entropy"
    weight: float = 1.0
    label_smoothing: float = 0.0
    reduction: str = "mean"
    custom_loss_fn: Optional[Callable] = None

    def get_loss_fn(self):
        """返回实际损失函数对象或函数"""
        if self.custom_loss_fn:
            return self.custom_loss_fn

        loss_type = (self.loss_type or "").lower()
        if loss_type == "cross_entropy":
            return nn.CrossEntropyLoss(label_smoothing=self.label_smoothing, reduction=self.reduction)
        elif loss_type == "mse":
            return nn.MSELoss(reduction=self.reduction)
        elif loss_type == "l1":
            return nn.L1Loss(reduction=self.reduction)
        elif loss_type == "focal":
            class FocalLoss(nn.Module):
                def __init__(self, gamma=2.0, reduction="mean"):
                    super().__init__()
                    self.gamma = gamma
                    self.reduction = reduction
                def forward(self, logits, target):
                    ce = nn.functional.cross_entropy(logits, target, reduction="none")
                    pt = torch.exp(-ce)
                    loss = ((1 - pt) ** self.gamma) * ce
                    if self.reduction == "mean":
                        return loss.mean()
                    elif self.reduction == "sum":
                        return loss.sum()
                    return loss
            return FocalLoss(reduction=self.reduction)
        else:
            raise ValueError(f"Unsupported loss type: {self.loss_type}")


# -----------------------------
# 多任务损失配置
# -----------------------------
@dataclass
class LossConfig:
    """
    支持多任务损失配置。
    losses: List[SingleLossConfig]
    """
    losses: List[SingleLossConfig] = field(default_factory=lambda: [SingleLossConfig()])

    def get_combined_loss(self):
        """
        返回一个函数，输入(pred_dict, target_dict)返回加权组合损失
        pred_dict / target_dict 可以是 dict 或 list
        """
        def combined_loss(pred_dict, target_dict):
            total_loss = 0.0
            for i, loss_cfg in enumerate(self.losses):
                loss_fn = loss_cfg.get_loss_fn()
                task_name = list(target_dict.keys())[i] if isinstance(target_dict, dict) else str(i)
                y_pred = pred_dict[task_name] if isinstance(pred_dict, dict) else pred_dict[i]
                y_true = target_dict[task_name] if isinstance(target_dict, dict) else target_dict[i]
                l = loss_fn(y_pred, y_true)
                total_loss += loss_cfg.weight * l
            return total_loss
        return combined_loss

    def summary(self):
        print("📌 多任务损失配置：")
        for i, loss_cfg in enumerate(self.losses):
            print(f"  Loss {i}: type={loss_cfg.loss_type}, weight={loss_cfg.weight}, label_smoothing={loss_cfg.label_smoothing}")


# -----------------------------
# 测试代码
# -----------------------------
if __name__ == "__main__":
    print("=== 测试 LossConfig ===")
    loss_config = LossConfig(
        losses=[
            SingleLossConfig(loss_type="cross_entropy", weight=1.0),
            SingleLossConfig(loss_type="mse", weight=0.5)
        ]
    )
    loss_config.summary()

    # 测试多任务损失计算
    combined_loss_fn = loss_config.get_combined_loss()
    logits_cls = torch.randn(4, 10)
    labels_cls = torch.randint(0, 10, (4,))
    pred_reg = torch.randn(4, 1)
    labels_reg = torch.randn(4, 1)

    pred_dict = {"cls": logits_cls, "reg": pred_reg}
    target_dict = {"cls": labels_cls, "reg": labels_reg}

    loss_value = combined_loss_fn(pred_dict, target_dict)
    print(f"计算的总损失: {loss_value.item():.4f}")
